import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt

# 指定数据集下载路径
data_dir = "E:\\tensorflow\\tensorflow_learn\\datasets"

# 加载 MNIST 数据集，并指定数据存储路径
# tfds.load: 加载指定的数据集，with_info=True 返回数据集信息，as_supervised=True 以 (image, label) 的形式返回数据
dataset, info = tfds.load('mnist', with_info=True, as_supervised=True, data_dir=data_dir)

# 获取训练集和测试集
train_dataset, test_dataset = dataset['train'], dataset['test']


# 数据预处理函数
# 定义预处理函数，将图像数据归一化到 [0, 1] 范围，并返回处理后的图像和标签
def preprocess(image, label):
    image = tf.cast(image, tf.float32) / 255.0
    return image, label


# 应用预处理并设置批大小
# 使用 map 方法对数据集中的每个元素应用预处理函数，然后分批并打乱训练集
BATCH_SIZE = 32
train_dataset = train_dataset.map(preprocess).batch(BATCH_SIZE).shuffle(1000)
test_dataset = test_dataset.map(preprocess).batch(BATCH_SIZE)

# 构建模型
# 使用 tf.keras.layers.Input 定义输入形状
model = tf.keras.models.Sequential([
    tf.keras.layers.Input(shape=(28, 28)),  # 定义输入形状
    tf.keras.layers.Flatten(),  # 将输入的 28x28 图像展平为一维数组
    tf.keras.layers.Dense(128, activation='relu'),  # 添加一个全连接层，128 个神经元，激活函数为 ReLU
    tf.keras.layers.Dropout(0.2),  # 添加 Dropout 层，随机丢弃 20% 的神经元，防止过拟合
    tf.keras.layers.Dense(10)  # 输出层，10 个神经元，对应 10 个类别
])

# 定义损失函数
# 使用 SparseCategoricalCrossentropy 损失函数，适用于多分类问题，from_logits=True 表示模型输出 logits
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 编译模型
# 使用 Adam 优化器，损失函数为 loss_fn，评估指标为准确率
model.compile(optimizer='adam', loss=loss_fn, metrics=['accuracy'])

# 训练模型
# 使用训练集训练模型，训练 5 轮
model.fit(train_dataset, epochs=5)

# 评估模型
# 使用测试集评估模型性能
model.evaluate(test_dataset)

# 将 logits 转换为概率
# 构建一个新的模型，将 logits 转换为概率
probability_model = tf.keras.Sequential([
    model,
    tf.keras.layers.Softmax()
])

# 从测试集中取出一些图像进行预测
# 从测试集中取出一个批次的图像和标签，使用封装后的模型进行预测
# 从测试集中取出一个批次的图像和标签
images, labels = next(iter(test_dataset.take(1)))

# 使用封装后的模型进行预测
predictions = probability_model(images)


# 可视化预测结果
# 定义可视化函数，显示图像及其预测结果
def plot_image(i, predictions_array, true_label, img):
    """

    :param i:
    :param predictions_array: 是 predictions[i]，表示第 i 个图像的预测概率数组。
    :param true_label:
    :param img:
    :return:
    """
    true_label, img = true_label[i], img[i]
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(img, cmap=plt.cm.binary)
    predicted_label = np.argmax(predictions_array)
    if predicted_label == true_label:
        color = 'blue'
    else:
        color = 'red'
    plt.xlabel("{} {:2.0f}% ({})".format(predicted_label,
                                         100 * np.max(predictions_array),
                                         true_label),
               color=color)


# 显示前 10 个图像的预测结果
# 创建一个大小为 10x5 的图形，显示前 10 个图像及其预测结果
plt.figure(figsize=(10, 5))
for i in range(10):
    plt.subplot(2, 5, i + 1)
    plot_image(i, predictions[i], labels, images)
plt.tight_layout()
plt.show()
